from typing import TypeVar

from zkl_ptutils_training import ModelFactory, SimpleTrainingEngine, SimpleTrainingEngineConfig, \
    SimpleTrainingEngineHparams

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')

class NbpttTrainingEngine(SimpleTrainingEngine):
    def __init__(self, *,
        model_factory: ModelFactory,
        config: SimpleTrainingEngineConfig | None = None,
        hparams: SimpleTrainingEngineHparams | None = None,
    ):
        super().__init__(
            model_factory=model_factory,
            config=config, hparams=hparams)
        self._policy_hist_grads = {}

    def backward(self, output: AnyOutput):
        gradients = output['gradients']
        for name, grad in gradients.items():
            parameter = self.model.model.get_parameter(name)
            if parameter.grad is None:
                parameter.grad = grad
            else:
                parameter.grad += grad
        self.accumulated_steps_n += 1
